import { AIProvider, AIProviderConfig, ImageGenerationRequest, ImageGenerationResponse } from './AIProvider';
import axios from 'axios';

export interface WanxConfig extends AIProviderConfig {
  model?: string;
  endpoint?: string;
  pollInterval?: number;
  maxPollAttempts?: number;
}

export class WanxProvider extends AIProvider {
  private wanxConfig: WanxConfig;
  private readonly defaultEndpoint = 'https://dashscope.aliyuncs.com/api/v1/services/aigc/text2image/image-synthesis';
  private readonly defaultModel = 'wanx2.1-imageedit';

  constructor(config: WanxConfig) {
    super(config);
    this.wanxConfig = {
      ...config,
      endpoint: config.endpoint || this.defaultEndpoint,
      model: config.model || this.defaultModel,
      pollInterval: config.pollInterval || 2000, // 2秒轮询间隔
      maxPollAttempts: config.maxPollAttempts || 60 // 最多轮询60次（2分钟）
    };
    this.validateConfig();
  }

  protected validateConfig(): void {
    if (!this.config.apiKey) {
      throw new Error('DashScope API Key is required for Wanx provider');
    }
  }

  isAvailable(): boolean {
    return !!this.config.apiKey;
  }

  async generateImage(request: ImageGenerationRequest): Promise<ImageGenerationResponse> {
    if (!this.isAvailable()) {
      throw new Error('Wanx provider is not available');
    }

    try {
      // 第一步：创建任务
      const taskId = await this.createTask(request);

      // 第二步：轮询获取结果
      const result = await this.pollTaskResult(taskId);

      return {
        imageUrl: result.imageUrl,
        requestId: taskId,
        metadata: {
          model: this.wanxConfig.model,
          taskId,
          provider: 'wanx'
        }
      };
    } catch (error) {
      console.error('Wanx image generation failed:', error);
      throw new Error(`Wanx image generation failed: ${error instanceof Error ? error.message : 'Unknown error'}`);
    }
  }

  private async createTask(request: ImageGenerationRequest): Promise<string> {
    const payload = {
      model: this.wanxConfig.model,
      input: {
        prompt: request.prompt,
        function: 'doodle', 
        base_image_url: request.originalImageUrl
      },
      parameters: {
        size: this.mapSize(request.size || '1024*1024'),
        n: request.n || 1,
        seed: Math.floor(Math.random() * 1000000)
      }
    };

    const response = await axios.post(this.wanxConfig.endpoint!, payload, {
      headers: {
        'Authorization': `Bearer ${this.config.apiKey}`,
        'Content-Type': 'application/json',
        'X-DashScope-Async': 'enable'
      },
      timeout: this.config.timeout || 30000
    });

    if (response.data.output?.task_id) {
      return response.data.output.task_id;
    }

    throw new Error(`Failed to create Wanx task: ${JSON.stringify(response.data)}`);
  }

  private async pollTaskResult(taskId: string): Promise<{ imageUrl: string }> {
    const pollUrl = `https://dashscope.aliyuncs.com/api/v1/tasks/${taskId}`;
    let attempts = 0;

    while (attempts < this.wanxConfig.maxPollAttempts!) {
      try {
        const response = await axios.get(pollUrl, {
          headers: {
            'Authorization': `Bearer ${this.config.apiKey}`
          },
          timeout: this.config.timeout || 10000
        });

        const {  output } = response.data;
        const { task_status } = output;

        if (task_status === 'SUCCEEDED') {
          if (output?.results?.[0]?.url) {
            return { imageUrl: output.results[0].url };
          }
          throw new Error('Task succeeded but no image URL found');
        }

        if (task_status === 'FAILED') {
          const errorMessage = output?.message || 'Task failed without error message';
          throw new Error(`Wanx task failed: ${errorMessage}`);
        }

        // 任务还在进行中，继续轮询
        if (task_status === 'PENDING' || task_status === 'RUNNING') {
          await this.sleep(this.wanxConfig.pollInterval!);
          attempts++;
          continue;
        }

        throw new Error(`Unknown task status: ${task_status}`);
      } catch (error) {
        if (axios.isAxiosError(error)) {
          console.error(`Polling attempt ${attempts + 1} failed:`, error.message);
        }

        if (attempts >= this.wanxConfig.maxPollAttempts! - 1) {
          throw error;
        }

        await this.sleep(this.wanxConfig.pollInterval!);
        attempts++;
      }
    }

    throw new Error(`Task polling timeout after ${attempts} attempts`);
  }

  private mapSize(size: string): string {
    // 将OpenAI格式的尺寸转换为万象格式
    const sizeMap: Record<string, string> = {
      '1024x1024': '1024*1024',
      '1024x1792': '1024*1792',
      '1792x1024': '1792*1024',
      '512x512': '512*512',
      '768x768': '768*768'
    };

    return sizeMap[size] || '1024*1024';
  }

  private sleep(ms: number): Promise<void> {
    return new Promise(resolve => setTimeout(resolve, ms));
  }

  getProviderName(): string {
    return 'Wanx';
  }
}